package com.ikas.ai.framework.node.impl;

import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.ikas.ai.framework.core.model.CalculateParamPair;
import com.ikas.ai.framework.core.model.NodeInputVariables;
import com.ikas.ai.framework.flow.ScoreFlow;
import com.ikas.ai.framework.jexl.JexlEngineFactory;
import com.ikas.ai.framework.node.NodeType;
import org.apache.commons.jexl3.JexlContext;
import org.apache.commons.jexl3.JexlEngine;
import org.apache.commons.jexl3.JexlScript;
import org.apache.commons.jexl3.MapContext;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @Description: 数学公式计算
 * @Author: qi.yilai
 * @Date: 2023/1/12 17:33
 * @Version: 1.0.0
 **/
public class MathCalculateNode extends AbstractNode<JSONObject> {
    private static final String VARIABLE_PREFIX = "var";
    private static final Pattern VARIABLE_PATTERN = Pattern.compile("`([^`]+)`");

    private List<CalculateParamPair> params;

    // 中间计算出来的临时变量
    private JSONObject calcVariables = new JSONObject();
    private AtomicLong atomicLong = new AtomicLong(1);

    public MathCalculateNode(ScoreFlow scoreFlow, List<CalculateParamPair> params, Long nodeId, Long nodeInstanceId, String name) {
        super(scoreFlow, nodeId, nodeInstanceId, name);
        this.params = params;
    }

    @Override
    public NodeType type() {
        return NodeType.MATH_CALC;
    }

    @Override
    JSONObject doExecute() {
        JSONObject result = new JSONObject();
        params.forEach(calculateParamPair -> {
            JexlEngine engine = JexlEngineFactory.getInstance();

            JexlScript script;
            JexlContext jexlContext;
            if (needReplace(calculateParamPair.getExpression())) {
                Set<String> variables = findVariableForReplace(calculateParamPair.getExpression());
                Map<String, String> renameMap = rename(previous, variables);
                script = engine.createScript(replaceExpression(calculateParamPair.getExpression(), renameMap));
                jexlContext = createJexlContext(previous, renameMap);
            } else {
                script = engine.createScript(calculateParamPair.getExpression());
                jexlContext = new MapContext();
                for (Map.Entry<String, Object> entry : previous.getParams().entrySet()) {
                    jexlContext.set(entry.getKey(), entry.getValue());
                }
            }

            Object value = script.execute(jexlContext);
            String resultVariable = getOriginVariable(calculateParamPair.getResultVariable());
            result.put(resultVariable, value);
            calcVariables.put(resultVariable, value);
        });

        return result;
    }

    private String getOriginVariable(String variableName) {
        if (variableName.startsWith("`") && variableName.endsWith("`")) {
            return variableName.substring(1, variableName.length() -1);
        } else {
            return variableName;
        }
    }

    private Map<String, String> rename(NodeInputVariables nodeInputVariables, Set<String> variables) {
        Map<String, String> renameMap = Maps.newHashMap();
        JSONObject params = nodeInputVariables.getParams();
        for (String var : variables) {
            String originalVar = getOriginVariable(var);
            if (params.containsKey(originalVar)) {
                String rename = VARIABLE_PREFIX + atomicLong.getAndIncrement();
                renameMap.put(originalVar, rename);
            } else {
                throw new IllegalArgumentException(String.format("表达式变量【%s】不存在", originalVar));
            }
        }
        return renameMap;
    }

    private boolean needReplace(String expression) {
        Matcher matcher = VARIABLE_PATTERN.matcher(expression);
        return matcher.find();
    }

    private Set<String> findVariableForReplace(String expression) {
        Set<String> variables = Sets.newHashSet();
        Matcher matcher = VARIABLE_PATTERN.matcher(expression);
        while (matcher.find()) {
            variables.add(matcher.group());
        }
        return variables;
    }

    private String replaceExpression(String expression, Map<String, String> variableMap) {
        Matcher matcher = VARIABLE_PATTERN.matcher(expression);
        String replaceExpr = expression;
        while (matcher.find()) {
            String rename = variableMap.get(matcher.group(1));
            replaceExpr = matcher.replaceFirst(rename);
            matcher = VARIABLE_PATTERN.matcher(replaceExpr);
        }
        return replaceExpr;
    }

    private JexlContext createJexlContext(NodeInputVariables nodeInputVariables, Map<String, String> renameMap) {
        JexlContext jexlContext = new MapContext();
        for (Map.Entry<String, Object> entry : nodeInputVariables.getParams().entrySet()) {
            if (renameMap.containsKey(entry.getKey())) {
                jexlContext.set(renameMap.get(entry.getKey()), entry.getValue());
            } else {
                jexlContext.set(entry.getKey(), entry.getValue());
            }
        }
        return jexlContext;
    }

    public JSONObject calcVariables() {
        return calcVariables;
    }
}
